"""
手写体数字识别模型的训练和保存
数据来源：MNIST数据库（tensorflow.examples.tutorials.mnist）
        Modified National Institute of Standards and Technology database
        修改后的美国国家标准与技术研究所数据库
训练集：mnist.train
测试集：mnist.test
每个样本：28*28=784的灰度图片
神经网络结构：输入层(784个输入) → 隐藏层(10个神经元+softmax激活) → 输出层
"""
import os
import cv2
import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data


class HandwrittenNumNN:
    def __init__(self, data_dir='', model_dir=None):
        """
        构造方法
        :param data_dir: 数据集所在目录，默认从网络下载
        :param model_dir: 模型的存储目录，为None时则不保存模型
        """
        if type(model_dir) == str:
            if not model_dir.endswith('/'):
                raise Exception('model_dir必须以“/”结尾')

        self.current_dir = os.path.dirname(__file__)
        self.model_dir = model_dir and os.path.join(self.current_dir, model_dir)
        self.data_dir = os.path.join(self.current_dir, data_dir)
        self.sess = tf.Session()
        self.data = input_data.read_data_sets(self.data_dir, one_hot=True)
        # print(self.data.train.next_batch(1))

        # 原始输入(真实数据)
        self.input = tf.placeholder(tf.float32, [None, 784])
        # print(self.input)  # shape=(?, 784)

        # 原始输入对应的标签(真实结果)
        self.label = tf.placeholder(tf.float32, [None, 10])

        # 神经网络模型(结果)
        self.nn = self.create_nn()

        # 模型准确率
        self.accuracy = self.calc_accuracy()

        # 训练步骤
        self.train_step = self.create_train_step()

        # 模型状态保存器
        self.saver = tf.train.Saver()

        # 模型变量初始化
        self.init_vars()

    def create_nn(self):
        """
        创建神经网络
        :return: 返回网络输出结果
        """
        # 输入层(接全连接层)
        input = self.input

        # 全连接层 输出10个神经元 激活 最终结果
        fc = self.create_fc(input, 784, 10)
        # print(fc)  # shape=(?, 10)

        return fc

    def create_fc(self, input, input_features, out_features, active=True):
        """
        创建全连接层
        :param input: 输入（经过拉伸的一维张量）
        :param input_features: 输入特征值数量
        :param out_features: 输出特征值数量（神经元数量）
        :param active: 是否激活
        :return:
        """
        # 初始化权重
        weight = self.init_weight([input_features, out_features])

        # 初始化偏置，每个输出(神经元)对应一个偏置和n个权重
        bias = self.init_bias([out_features])

        # 神经网络计算
        fc = tf.matmul(input, weight) + bias

        # 激活运算
        if active:
            fc = tf.nn.softmax(fc)

        return fc

    def init_weight(self, shape):
        """
        初始化权重(卷积核可以理解为特殊的权重)
        :param shape: 权重张量形状
        :return: 初始化后的权重值
        """
        initial = tf.random_normal(shape)
        return tf.Variable(initial)

    def init_bias(self, shape):
        """
        初始化偏置
        :param shape: 形状
        :return: 初始化后的偏置值
        """
        initial = tf.zeros(shape)
        return tf.Variable(initial)

    def calc_accuracy(self):
        """
        计算模型准确率
        :return: 返回模型准确率
        """
        # 评价
        correct_prediction = tf.equal(
            tf.argmax(self.nn, 1),  # 取出预测概率中最大的值的索引
            tf.argmax(self.label, 1)  # 取出真实概率中最大的值的索引
        )

        # 将上一步得到的bool类型数组转换为浮点型，并求准确率
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

        return accuracy

    def create_train_step(self):
        """
        创建训练操作
        :return: 返回训练操作
        """
        # 损失函数
        loss_func = -tf.reduce_sum(self.label * tf.log(self.nn), 1)

        # 交叉熵
        cross_entropy = tf.reduce_mean(loss_func)

        # 优化器
        optimizer = tf.train.GradientDescentOptimizer(0.0001)

        # 步骤
        train_step = optimizer.minimize(cross_entropy)

        return train_step

    def init_vars(self, restore=True):
        """
        模型全局变量初始化
        :param restore: 是否从保存在模型目中的checkpoint恢复全局变量
        :return:
        """
        self.sess.run(tf.global_variables_initializer())
        if restore and self.model_dir:
            checkpoint = os.path.join(self.model_dir, 'checkpoint')
            if os.path.exists(checkpoint):
                self.saver.restore(self.sess, self.model_dir)
            else:
                print('模型文件路径有误')

    def save_vars(self):
        """
        保存模型变量
        :return:
        """
        if self.model_dir:
            self.saver.save(self.sess, self.model_dir)
            print('save model success')

    def train(self):
        """
        执行训练操作
        :return:
        """
        epochs = 200  # 训练轮次
        batch_size = 100  # 批次大小

        for i in range(epochs):  # 迭代训练
            total_batch = int(self.data.train.num_examples / batch_size)  # 计算批次数量，int为下取整
            for j in range(total_batch):
                batch = self.data.train.next_batch(batch_size)  # 获取一个批次样本
                params = {
                    self.input: batch[0],  # 输入
                    self.label: batch[1],  # 标签
                }

                t, acc = self.sess.run(
                    [self.train_step, self.accuracy],  # op
                    params  # 喂入参数
                )

                if j % 100 == 0:
                    print("epoch: %d, pass: %d, acc: %f" % (i, j, acc))

    def eval(self, input, label):
        """
        模型测试评估
        :param input: 输入（测试数据）
        :param label: 标签（真实结果）
        :return: 返回测试准确率
        """
        params = {self.input: input, self.label: label}
        test_acc = self.sess.run(self.accuracy, params)
        print('Test accuracy %f' % test_acc)
        return test_acc

    def close(self):
        self.sess.close()

    def recognize(self, image):
        """
        识别一张数字手写体图片（白底黑字）
        :param image: 能被cv2识别的图片数据
        :return: 识别的结果
        """
        img = image
        if len(image.shape) == 3:
            img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        img = cv2.resize(img, (28, 28))
        img = 1 - img / 255
        # cv2.imshow('img',img)
        # cv2.waitKey()
        # cv2.destroyAllWindows()
        params = {self.input: img.reshape((1, 784))}
        result = self.sess.run(self.nn, params)
        result = np.argmax(result[0])
        return result


handwritten_num_nn = HandwrittenNumNN(
    data_dir='MNIST_data/handwritten_num/',
    model_dir='AI_model/handwritten_num/'
)

if __name__ == '__main__':
    nn = handwritten_num_nn

    # print("begin training...")
    # nn.train()
    # nn.save_vars()

    print('\nbegin testing...')
    nn.eval(nn.data.test.images, nn.data.test.labels)
    #
    # while True:
    #     img_path = input('请输入图片路径：')
    #     if not img_path:
    #         break
    #     img = cv2.imread(img_path, 0)
    #     cv2.imshow(img_path, img)
    #     result = nn.recognize(img)
    #     print('该图片的结果：', result)
    #     cv2.waitKey()
    #     cv2.destroyAllWindows()

    nn.close()
